package com.feidee.fd.sml.algorithm.component.ml.classification

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.NaiveBayes

/**
  * @Author songhaicheng
  * @Date 2018/08/28
  * @Email: haicheng_song@sui.com
  */
case class NaiveBayesParam(
                            override val input_pt: String,
                            override val output_pt: String,
                            override val hive_table: String,
                            override val flow_time: String,
                            override val featuresCol: String,
                            override var labelCol: String,
                            override var predictionCol: String,
                            override val modelPath: String,
                            override val metrics: Array[String],
                            // 原始预测值列名，默认 rawPrediction
                            rawPredictionCol: String,
                            // 预测概率列名，默认 probability
                            probabilityCol: String,
                            // 样本权重列
                            weightCol: String,
                            // 平滑参数，需大于或者等于 0.0，默认 1.0
                            smoothing: Double,
                            // 模型类型，仅支持 multinomial（算法默认）& bernoulli
                            modelType: String,
                            // 多分类的类别判定概率的阈值，数量需与分类树一致
                            thresholds: Array[Double]
                          ) extends MLParam {

  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    "rawPrediction", "probability", null, 1.0, null, new Array[Double](0))

  override def verify(): Unit = {
    super.verify()
    require(smoothing >= 0.0, "param smoothing can't be negative")
    require(!tool.isNotNull(modelType) ||
      (tool.isNotNull(modelType) && Array("multinomial", "bernoulli").contains(modelType.toLowerCase)),
      s"param modelType only accepts [multinomial, bernoulli], but has $modelType")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("probabilityCol" -> probabilityCol)
    map += ("weightCol" -> weightCol)
    map += ("smoothing" -> smoothing)
    map += ("modelType" -> modelType)
    map += ("thresholds" -> thresholds)
    map
  }
}


class NaiveBayesComponent extends AbstractClassificationComponent[NaiveBayesParam] {

  override def setUp(param: NaiveBayesParam): PipelineStage = {
    val nb = new NaiveBayes()
      .setLabelCol(param.labelCol)
      .setFeaturesCol(param.featuresCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setPredictionCol(param.predictionCol)
      .setProbabilityCol(param.probabilityCol)

    if (tool.isNotNull(param.weightCol)) {
      nb.setWeightCol(param.weightCol)
    }
    if (param.smoothing >= 0.0) {
      nb.setSmoothing(param.smoothing)
    }
    if (tool.isNotNull(param.modelType)) {
      nb.setModelType(param.modelType.toLowerCase)
    }
    if (param.thresholds.length > 0) {
      nb.setThresholds(param.thresholds)
    }
    nb
  }

}

object NaiveBayesComponent {

  def apply(paramStr: String): Unit = {
    new NaiveBayesComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    NaiveBayesComponent(args(0))
  }

}
